import numpy as np
import matplotlib.pyplot as plt
import struct
from scipy.signal import butter,lfilter,freqz

# 低通滤波
# y(n)=2y(n-1)-y(n-2)+x(n)-2x(n-6)+x(n-12)
class LowPass():
    def __init__(self):
        self.Navg = 12 #输入N实现N阶的平滑
        self.xbuf = self.Navg*[0] #存储x(n-N)...x(n-1)
        self.ptr = 0 #指向最新的数据存放位置，x(n)
        self.yn = 0 #存储y(n)
        self.yn1 = 0
        self.yn2 = 0
        pass

    def lfilter(self, xn):
        n = self.ptr - self.Navg
        if n < 0:
            n += self.Navg
        self.yn = 2 * self.yn1 - self.yn2 + xn - 2 * self.xbuf[self.ptr - 6] + self.xbuf[n]
        self.xbuf[self.ptr] = xn
        self.ptr += 1
        if self.ptr == self.Navg:
            self.ptr = 0
        self.yn2 = self.yn1
        self.yn1 = self.yn
        return self.yn/36

# 高通滤波
#y(n)=32*x(n-16)-[y(n-1)+x(n)-x(n-32)]
#y(n)=y(n-1)+x(n)-x(n-32),z(n)=x(n-16)-y(n)/32;
class HighPass():

    def __init__(self):
        self.Navg = 32
        self.xbuf = self.Navg*[0]
        self.ptr = 0
        self.yn = 0
        self.yn1 = 0
        self.zn = 0

    def hfilter(self,xn):
        n = self.ptr - self.Navg
        if n < 0:
            n += self.Navg
        self.yn = self.yn1 + xn - self.xbuf[n]
        self.zn = self.xbuf[self.ptr - 16] - (self.yn / 32)
        self.xbuf[self.ptr] = xn
        self.ptr += 1
        if self.ptr == self.Navg:
            self.ptr = 0
        self.yn1 = self.yn
        return self.zn

#差分滤波
#y(n)=(1/8)*[-x(n-2)-2*x(n-1)+2x(n+1)+x(n+2)]
#y(n)=（1/8)*[-x(n-4)-2x(n-3)+2x(n-1)+x(n)]
class ecg():
    def __init__(self):
        self.Navg=4
        self.xbuf=self.Navg*[0]
        self.ptr = 0
        self.yn = 0

    def cffilter(self, xn):
        self.yn = -1*self.xbuf[self.ptr]-2*self.xbuf[self.ptr-3]+2*self.xbuf[self.ptr-1]+xn
        self.xbuf[self.ptr] = xn
        self.ptr = self.ptr+1
        if self.ptr == self.Navg:
            self.ptr = 0
        return self.yn/8
#平方
class squaring():
    def __init__(self):
        y = 0
    def pffilter(self,xn):
        yn = xn*xn
        return yn
#平滑
#y(n)=1/N[x(n-(N-1))+x(n-(N-2))+...+x(n)]
# N = 40
class Smooth():
    def __init__(self):
        self.Navg=40
        self.xbuf=self.Navg*[0]
        self.ptr=0
        self.yn=0

    def phfilter(self,xn):
        n = self.ptr - self.Navg
        if n < 0:
            n += self.Navg
        self.yn = self.yn+xn-self.xbuf[n]
        self.xbuf[self.ptr] = xn
        self.ptr += 1
        if self.ptr == self.Navg:
            self.ptr = 0
        return self.yn/self.Navg


class QRSFilter:
    def __init__(self):
        self.lp = LowPass()
        self.hp = HighPass()
        self.diff = ecg()
        self.squ = squaring()
        self.sm = Smooth()
        pass

    def filter(self,ecg):
        x = self.lp.lfilter(ecg)
        x1 = self.hp.hfilter(x)
        x2 = self.diff.cffilter(x1)
        x3 = self.squ.pffilter(x2)
        x4 = self.sm.phfilter(x3)
        return x4


if __name__ == '__main__':
    fp = open("D:/wuhui/testcode/SmartHealth-master/DATA/mitdb/102.dat","rb")
    fp_out = open("D:/wuhui/testcode/SmartHealth-master/DATA/mitdb/102_out.dat","wb")
    fp_out1 = open("D:/wuhui/testcode/SmartHealth-master/DATA/mitdb/102_out1.dat","wb")
    diff0 = LowPass()
    diff1 = HighPass()
    diff2 = ecg()
    diff3 = squaring()
    diff4 = Smooth()

    for i in range(0, 1000):
        x = struct.unpack('h', fp.read(2))
        y0 = diff0.lfilter(x[0])
        y1 = diff1.hfilter(y0)
        y2 = diff2.cffilter(y1)
        y3 = diff3.pffilter(y2)
        y4 = diff4.phfilter(y3)

        fp_out.write(struct.pack('h', int(y2))) #将字符串写入到打开的文件中
        fp_out1.write(struct.pack('h', int(y4)))

    fp.close()#close() 函数是专门用来关闭已打开文件的
    fp_out.close()
    fp_out1.close()
    ecg0 = np.fromfile("D:/wuhui/testcode/SmartHealth-master/DATA/mitdb/102.dat",dtype=np.short)#按照short格式读取数据
    ecg1 = np.fromfile("D:/wuhui/testcode/SmartHealth-master/DATA/mitdb/102_out.dat",dtype=np.short)
    ecg2 = np.fromfile("D:/wuhui/testcode/SmartHealth-master/DATA/mitdb/102_out1.dat",dtype=np.short)
    
    plt.subplot(3, 1, 1)
    plt.plot(ecg0[0:1000])
    plt.subplot(3, 1, 2)
    plt.plot(ecg1[0:1000])
    plt.subplot(3, 1, 3)
    plt.plot(ecg2[0:1000])
    plt.show()